#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：pyspark 
@File    ：MovieRatings.py
@IDE     ：PyCharm 
@Author  ：Amos.Mu
@Date    ：2022/5/19 21:48 
'''

import os
import pandas as pd
import numpy as np

DATA_PATH = "../ml-latest-small/ratings.csv"
CACHE_DIR = "../cache/"


def load_data(data_path):
    '''
    加载数据
    :param data_path: 数据文件路径
    :return: 用户物品评分矩阵
    '''
    # 数据缓存地址
    cache_path = os.path.join(CACHE_DIR, "ratings_matrix.cache")

    print("开始加载数据集...")
    if os.path.exists(cache_path):  # 如果缓存文件中有数据
        print("数据加载中...")
        ratings_matrix = pd.read_pickle(cache_path)
        print("从缓存中加载数据完毕")
    else:
        print("加载数据中...")
        dtype = {"userId": np.int32, "movieId": np.int32, "rating": np.float32}
        ratings = pd.read_csv(data_path, dtype=dtype, usecols=range(3))
        ratings_matrix = ratings.pivot_table(index=["userId"], columns=["movieId"], values="rating")
        ratings_matrix.to_pickle(cache_path)
    return ratings_matrix


def compute_pearson_similarity(ratings_matrix, based="user"):
    '''
    计算皮尔逊相关系数
    :param ratings_matrix: 评分矩阵
    :param based: user 或者 item
    :return: 相似度矩阵
    '''

    user_similarity_cacge_path = os.path.join(CACHE_DIR, "user_similarity.cache")
    item_similarity_cacge_path = os.path.join(CACHE_DIR, "item_similarity.cache")

    # 基于皮尔逊相关系数计算相似度
    # 1.用户相似度
    if based == "user":
        if os.path.exists(user_similarity_cacge_path):
            print("正在加载用户相似度矩阵")
            similarity = pd.read_pickle(user_similarity_cacge_path)
        else:
            print("开始计算用户相似度矩阵")
            similarity = ratings_matrix.T.corr()
            similarity.to_pickle(user_similarity_cacge_path)
    elif based == "item":
        if os.path.exists(user_similarity_cacge_path):
            print("正在加载用户相似度矩阵")
            similarity = pd.read_pickle(user_similarity_cacge_path)
        else:
            print("开始计算用户相似度矩阵")
            similarity = ratings_matrix.T.corr()
            similarity.to_pickle(user_similarity_cacge_path)
    else:
        raise Exception("Unhandled 'based' Value:%s" % based)
    print("相似度矩阵计算/加载完毕")
    return similarity


def predict(uid, iid, ratings_matrix, user_similar):
    '''
    预测给定用户对给定物品的评分值 (user-based)
    :param uid: 用户id
    :param iid: 物品id
    :param ratings_matrix: 用户-物品 评分矩阵
    :param user_similar: 用户两两相似矩阵
    :return: 预测的评分矩阵
    '''
    print("开始预测用户<%d>对电影<%d>的评分..." % (uid, iid))
    # 1.找到uid用户的相似用户  取出uid的相似用户向量 删除用户本身和空值
    similar_users = user_similar[uid].drop([uid]).dropna()

    # 相似用户筛选规则：正相关的用户
    similar_users = similar_users.where(similar_users > 0).dropna()
    if similar_users.empty is True:
        raise Exception("用户<%d>没有相似的用户" % uid)

    # 2.从uid用户的近邻相似用户中筛选出对iid物品有评分记录的近邻用户
    # &是两个set求交集
    ids = set(ratings_matrix[iid].dropna().index) & set(similar_users.index)
    finally_similar_users = similar_users.loc[list(ids)]
    print(finally_similar_users)

    # 3.结合uid用户与其近邻相似用户中筛选出对iid物品有评分记录的近邻用户
    '''
      这里是使用user-based(CF)进行评分预估
    '''
    sum_up = 0  # 评分预测公式的分子部分
    sum_down = 0  # 评分预测公式的分母部分
    for sim_uid, similarity in finally_similar_users.iteritems():
        # 近邻用户的评分数据
        sim_user_rated_movies = ratings_matrix.loc[sim_uid].dropna()
        # 近邻用户对iid物品的评分
        sim_user_rating_for_item = sim_user_rated_movies[iid]
        # 计算分子的值
        sum_up += similarity * sim_user_rating_for_item

        #     计算分母的值
        sum_down += similarity
    # 计算预测的评分值并返回
    predict_rating = sum_up / sum_down
    print("预测出用户<%d>对电影<%d>的评分：%0.2f" % (uid, iid, predict_rating))
    return round(predict_rating, 2)


def predict_all(uid, ratings_matrix, user_similar):
    '''

    :param uid:
    :param item_ids:
    :param ratings_matrix:
    :param user_similar:
    :return:
    '''
    item_ids = ratings_matrix.columns
    for iid in item_ids:
        try:
            rating = predict(uid, iid, ratings_matrix, user_similar)
        except Exception as e:
            print(e)
        else:
            yield uid, iid, rating


def _predict_all(uid, item_ids, ratings_matrix, user_similar):
    '''
    预测全部评分
    :param uid: 用户id
    :param item_ids: 要预测的物品id列表
    :param ratings_matrix: 用户-物品打分矩阵
    :param user_similar: 用户两两间的相似度
    :return: 生成器，逐个返回预测评分
    '''

    # 逐个预测
    for iid in item_ids:
        try:
            rating = predict(uid, iid, ratings_matrix, user_similar)
        except Exception as e:
            print(e)
        else:
            yield uid, iid, rating


def predict_all_filter(uid, ratings_matrix, user_similar, filter_rule=None):
    '''
    预测全部评分，并可根据条件进行前置过滤
    :param uid: 用户ID
    :param ratings_matrix: 用户-物品 评分矩阵
    :param user_similar: 用户相似度矩阵
    :param filter_rule: 过滤规则
    :return: 生成器 逐个返回预测评分
    '''
    if not filter_rule:
        item_ids = ratings_matrix.columns()
    elif isinstance(filter_rule, str) and filter_rule == "unhot":
        '''
        过滤非热门电影
        '''
        count = ratings_matrix.count()
        '''过滤出评分高于10的电影 作为热门电影'''
        item_ids = count.where(count > 10).dropna().index
    elif isinstance(filter_rule, str) and filter_rule == "rated":
        '''过滤用户评分过的电影'''
        user_ratings = ratings_matrix.loc[uid]
        ratings = user_ratings < 6
        item_ids = ratings.where(ratings == False).dropna().index
    elif isinstance(filter_rule, list) and set(filter_rule) == set(["rated", "unhot"]):
        count = ratings_matrix.count()
        '''过滤出评分高于10的电影 作为热门电影'''
        ids1 = count.where(count > 10).dropna().index

        user_ratings = ratings_matrix.loc[uid]
        ratings = user_ratings < 6
        ids2 = ratings.where(ratings == False).dropna().index
        # 取二者交集
        item_ids = set(ids1) & set(ids2)
    else:
        raise Exception("无效的过滤参数")
    yield from _predict_all(uid, item_ids, ratings_matrix, user_similar)


if __name__ == '__main__':
    pd.options.display.max_rows = 610  # 显示610行
    ratings_matrix = load_data(DATA_PATH)
    print(ratings_matrix)
    similarity = compute_pearson_similarity(ratings_matrix)
    print(similarity)
    for result in predict_all(1, ratings_matrix, similarity):
        print(result)

    # for result in predict_all_filter(1, ratings_matrix, user_similar=similarity, filter_rule=["unhot", "rated"]):
    #     print(result)
